#!/usr/bin/env python3
import rospy
from sensor_msgs.msg import JointState
from geometry_msgs.msg import PoseStamped
from piper_sdk import C_PiperInterface
import numpy as np
from tf.transformations import quaternion_from_euler
import tf2_ros
import geometry_msgs.msg

class PiperArmNode:
    def __init__(self):
        # 初始化ROS节点
        rospy.init_node('piper_arm_node_single', anonymous=False)
        # rospy.init_node('piper_arm_node')

        # 初始化左臂的发布者
        self.joint_pub = rospy.Publisher('arm/joint_states', JointState, queue_size=10)
        self.pose_pub = rospy.Publisher('arm/end_pose', PoseStamped, queue_size=10)
        self.end_rotation_pub = rospy.Publisher('arm/end_rotation_raw', PoseStamped, queue_size=10)
        
        # 连接硬件
        self.pi_per = C_PiperInterface("can0")  # 修改为对应您设备的总线名称
        self.pi_per.ConnectPort()

        # 设置循环频率（约200Hz）
        self.rate = rospy.Rate(200)

         # 创建TF广播器
        self.tf_broadcaster = tf2_ros.TransformBroadcaster()

    def convert_joint_state(self, raw_data, is_gripper=False):
        """转换关节状态数据"""
        if is_gripper:
            return round((raw_data * 0.001) / 70.0, 8)  # 夹具转换公式
        return round(raw_data * 0.001 / 57.3, 8)  # 关节转换公式

    def publish_arm_data(self):
        try:
            # 获取原始数据
            joint_data = self.pi_per.GetArmJointMsgs()
            pose_data = self.pi_per.GetArmEndPoseMsgs()
            gripper_data = self.pi_per.GetArmGripperMsgs().gripper_state.grippers_angle
        except Exception as e:
            rospy.logerr(f"Failed to get arm data: {str(e)}")
            return

        # 创建JointState消息
        joint_msg = JointState()
        joint_msg.header.stamp = rospy.Time.now()
        joint_msg.name = ['joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6', 'gripper']
        joint_msg.position = [
            self.convert_joint_state(joint_data.joint_state.joint_1),
            self.convert_joint_state(joint_data.joint_state.joint_2),
            self.convert_joint_state(joint_data.joint_state.joint_3),
            self.convert_joint_state(joint_data.joint_state.joint_4),
            self.convert_joint_state(joint_data.joint_state.joint_5),
            self.convert_joint_state(joint_data.joint_state.joint_6),
            self.convert_joint_state(gripper_data, is_gripper=True)
        ]

        # 创建包含位置和旋转的PoseStamped消息
        pose_msg = PoseStamped()
        pose_msg.header.stamp = joint_msg.header.stamp
        pose_msg.header.frame_id = "base_link"
        
        # 设置位置
        pose_msg.pose.position.x = pose_data.end_pose.X_axis /1000.0/1000.0 # m
        pose_msg.pose.position.y = pose_data.end_pose.Y_axis /1000.0/1000.0
        pose_msg.pose.position.z = pose_data.end_pose.Z_axis /1000.0/1000.0
        
        # 将欧拉角转换为四元数 (假设原始数据是度)
        roll_rad = np.radians(pose_data.end_pose.RX_axis /1000.0)
        pitch_rad = np.radians(pose_data.end_pose.RY_axis /1000.0)
        yaw_rad = np.radians(pose_data.end_pose.RZ_axis /1000.0)

        quaternion = quaternion_from_euler(roll_rad, pitch_rad, yaw_rad, axes='sxyz')
    
        pose_msg.pose.orientation.x = quaternion[0]
        pose_msg.pose.orientation.y = quaternion[1]
        pose_msg.pose.orientation.z = quaternion[2]
        pose_msg.pose.orientation.w = quaternion[3]

        # 发布 base_link 到 arm_end_effector 的变换
        transformStamped = geometry_msgs.msg.TransformStamped()
        transformStamped.header.stamp = rospy.Time.now()
        transformStamped.header.frame_id = "base_link"
        transformStamped.child_frame_id = "arm_end_effector"
        
        transformStamped.transform.translation.x = pose_msg.pose.position.x
        transformStamped.transform.translation.y = pose_msg.pose.position.y
        transformStamped.transform.translation.z = pose_msg.pose.position.z
        transformStamped.transform.rotation = pose_msg.pose.orientation
        
        self.tf_broadcaster.sendTransform(transformStamped)

        # end_rotation_msg = PoseStamped()
        # end_rotation_msg.header.stamp = joint_msg.header.stamp
        # end_rotation_msg.header.frame_id = "base_link"
        # end_rotation_msg.pose.position.x = pose_data.end_pose.RX_axis
        # end_rotation_msg.pose.position.y = pose_data.end_pose.RY_axis
        # end_rotation_msg.pose.position.z = pose_data.end_pose.RZ_axis
        # self.end_rotation_pub.publish(end_rotation_msg)

        # 发布消息
        self.joint_pub.publish(joint_msg)
        self.pose_pub.publish(pose_msg)


    def run(self):
        while not rospy.is_shutdown():
            try:
                self.publish_arm_data()
                self.rate.sleep()
            except Exception as e:
                rospy.logerr(f"Error: {str(e)}")
                break

        # 关闭连接
        try:
            self.pi_per.DisconnectPort()
        except Exception as e:
            rospy.logerr(f"Error disconnecting: {str(e)}")

if __name__ == '__main__':
    try:
        node = PiperArmNode()
        node.run()
    except rospy.ROSInterruptException:
        rospy.loginfo("Node interrupted")
    except Exception as e:
        rospy.logerr(f"Unhandled exception: {str(e)}")
